{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# YOLOv5 Classification Tutorial\n",
        "\n",
        "YOLOv5 supports classification tasks too. This is the official YOLOv5 classification notebook tutorial. YOLOv5 is maintained by [Ultralytics](https://github.com/ultralytics/yolov5).\n",
        "\n",
        "This notebook covers:\n",
        "\n",
        "*   Inference with out-of-the-box YOLOv5 classification on ImageNet\n",
        "*  [Training YOLOv5 classification](https://blog.roboflow.com//train-YOLOv5-classification-custom-data) on custom data\n",
        "\n",
        "*Looking for custom data? Explore over 66M community datasets on [Roboflow Universe](https://universe.roboflow.com).*\n",
        "\n",
        "This notebook was created with Google Colab. [Click here](https://colab.research.google.com/drive/1FiSNz9f_nT8aFtDEU3iDAQKlPT8SCVni?usp=sharing) to run it."
      ],
      "metadata": {
        "id": "5GYQX3of4QiW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n",
        "\n",
        "Pull in respective libraries to prepare the notebook environment."
      ],
      "metadata": {
        "id": "-PJ8vlYXbWtN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# save the working directory path for later use\n",
        "\n",
        "import os\n",
        "HOME = os.getcwd()\n",
        "print(HOME)"
      ],
      "metadata": {
        "id": "6W8YIH8ft0D5",
        "outputId": "0f12c5fa-4e52-4f3e-cd31-84c048418e23",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/yolov5\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "pIM7fOwm8A7l",
        "outputId": "b6fc23a5-a582-40c3-f8c2-dec1efeee37c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Setup complete ✅ (12 CPUs, 83.5 GB RAM, 32.0/235.7 GB disk)\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/ultralytics/yolov5  # clone\n",
        "%cd yolov5\n",
        "%pip install -qr requirements.txt  # install\n",
        "\n",
        "import torch\n",
        "import utils\n",
        "display = utils.notebook_init()  # checks"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Infer on ImageNet\n",
        "\n",
        "To demonstrate YOLOv5 classification, we'll leverage an already trained model. In this case, we'll download the ImageNet trained models pretrained on ImageNet using YOLOv5 Utils."
      ],
      "metadata": {
        "id": "i_DrUi2nmF40"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from utils.downloads import attempt_download\n",
        "\n",
        "p5 = ['n', 's', 'm', 'l', 'x']  # P5 models\n",
        "cls = [f'{x}-cls' for x in p5]  # classification models\n",
        "\n",
        "for x in cls:\n",
        "    attempt_download(f'weights/yolov5{x}.pt')"
      ],
      "metadata": {
        "id": "o2scLEh6EYnL",
        "outputId": "292d9f5c-06ce-4dcf-f099-d834ae2aeba5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-cls.pt to weights/yolov5n-cls.pt...\n",
            "100%|██████████| 4.87M/4.87M [00:00<00:00, 59.3MB/s]\n",
            "\n",
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt to weights/yolov5s-cls.pt...\n",
            "100%|██████████| 10.5M/10.5M [00:00<00:00, 93.8MB/s]\n",
            "\n",
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m-cls.pt to weights/yolov5m-cls.pt...\n",
            "100%|██████████| 24.9M/24.9M [00:00<00:00, 112MB/s]\n",
            "\n",
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5l-cls.pt to weights/yolov5l-cls.pt...\n",
            "100%|██████████| 50.9M/50.9M [00:00<00:00, 91.3MB/s]\n",
            "\n",
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5x-cls.pt to weights/yolov5x-cls.pt...\n",
            "100%|██████████| 92.0M/92.0M [00:00<00:00, 133MB/s]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, we can infer on an example image from the ImageNet dataset."
      ],
      "metadata": {
        "id": "Fn2_a38DmZ2H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "PATH_TO_IMAGE = f\"{HOME}/bananas.jpg\"\n",
        "!curl https://i.imgur.com/OczPfaz.jpg -o {PATH_TO_IMAGE}"
      ],
      "metadata": {
        "id": "M8VGa2t8tpQm",
        "outputId": "938e5cd5-5f33-492a-a272-a3035652cc50",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100  408k  100  408k    0     0  2097k      0 --:--:-- --:--:-- --:--:-- 2107k\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Infer using classify/predict.py\n",
        "!python classify/predict.py --weights ./weigths/yolov5s-cls.pt --source bananas.jpg"
      ],
      "metadata": {
        "id": "qqxF5pHCrLd3",
        "outputId": "db59a5a0-0f27-48e8-94c0-88843c0634a3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[34m\u001b[1mclassify/predict: \u001b[0mweights=['./weigths/yolov5s-cls.pt'], source=bananas.jpg, data=data/coco128.yaml, imgsz=[224, 224], device=, view_img=False, save_txt=False, nosave=False, augment=False, visualize=False, update=False, project=runs/predict-cls, name=exp, exist_ok=False, half=False, dnn=False, vid_stride=1\n",
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n",
            "\n",
            "Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt to weigths/yolov5s-cls.pt...\n",
            "100% 10.5M/10.5M [00:00<00:00, 110MB/s]\n",
            "\n",
            "Fusing layers... \n",
            "Model summary: 117 layers, 5447688 parameters, 0 gradients, 11.4 GFLOPs\n",
            "image 1/1 /content/yolov5/bananas.jpg: 224x224 banana 0.96, zucchini 0.00, acorn squash 0.00, spaghetti squash 0.00, green mamba 0.00, 3.8ms\n",
            "Speed: 0.5ms pre-process, 3.8ms inference, 59.6ms NMS per image at shape (1, 3, 224, 224)\n",
            "Results saved to \u001b[1mruns/predict-cls/exp\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "From the output, we can see the ImageNet trained model correctly predicts the class `banana` with `0.95` confidence."
      ],
      "metadata": {
        "id": "yQmj7IXqo3kk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Train On Custom Data\n",
        "\n",
        "To train on custom data, we need to prepare a dataset with custom labels.\n",
        "\n",
        "To prepare custom data, we'll use [Roboflow](https://roboflow.com). Roboflow enables easy dataset prep with your team, including labeling, formatting into the right export format, deploying, and active learning with a `pip` package.\n",
        "\n",
        "If you need custom data, there are over 66M open source images from the community on [Roboflow Universe](https://universe.roboflow.com).\n",
        "\n",
        "(For more guidance, here's a detailed blog on [training YOLOv5 classification on custom data](https://blog.roboflow.com/train-YOLOv5-classification-custom-data).)\n",
        "\n",
        "\n",
        "Create a free Roboflow account, upload your data, and label.\n",
        "\n",
        "![](https://s4.gifyu.com/images/fruit-labeling.gif)"
      ],
      "metadata": {
        "id": "9bXHHYeVDCXg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load Custom Dataset\n",
        "\n",
        "Next, we'll export our dataset into the right directory structure for training YOLOv5 classification to load into this notebook. Select the `Export` button at the top of the version page, `Folder Structure` type, and `show download code`.\n",
        "\n",
        "The ensures all our directories are in the right format:\n",
        "\n",
        "```\n",
        "dataset\n",
        "├── train\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_123.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_456.jpg\n",
        "├── valid\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_789.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_101.jpg\n",
        "├── test\n",
        "│   ├── class-one\n",
        "│   │   ├── IMG_121.jpg\n",
        "│   └── class-two\n",
        "│       ├── IMG_341.jpg\n",
        "```\n",
        "\n",
        "![](https://i.imgur.com/BF9BNR8.gif)\n",
        "\n",
        "\n",
        "Copy and paste that snippet into the cell below."
      ],
      "metadata": {
        "id": "Cu6-lrukD6Hc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Ensure we're in the right directory to download our custom dataset\n",
        "import os\n",
        "os.makedirs(\"../datasets/\", exist_ok=True)\n",
        "%cd ../datasets/"
      ],
      "metadata": {
        "id": "6IIgJbP7G6Th",
        "outputId": "97d1fa28-b4d4-4be0-a27b-66d577172ee6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/datasets\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install roboflow==1.1.48 -q"
      ],
      "metadata": {
        "id": "gnmezvmAucvw",
        "outputId": "aa44dc75-ebf2-4291-e870-5df08abedc7d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/80.3 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m80.3/80.3 kB\u001b[0m \u001b[31m4.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m66.8/66.8 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.5/54.5 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# REPLACE the below with your exported code snippet from above\n",
        "import roboflow\n",
        "roboflow.login()\n",
        "\n",
        "rf = roboflow.Roboflow()\n",
        "project = rf.workspace(\"model-examples\").project(\"banana-ripeness-classification-3dyre\")\n",
        "dataset = project.version(1).download(\"folder\")"
      ],
      "metadata": {
        "id": "He6JwHIlG-W_",
        "outputId": "68bd7ec3-be6c-432b-cf6c-3566c74f2efe",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\rvisit https://app.roboflow.com/auth-cli to get your authentication token.\n",
            "Paste the authentication token here: ··········\n",
            "loading Roboflow workspace...\n",
            "loading Roboflow project...\n",
            "Exporting format folder in progress : 85.0%\n",
            "Version export complete for folder format\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading Dataset Version Zip in Banana-Ripeness-Classification-1 to folder:: 100%|██████████| 152300/152300 [00:02<00:00, 73930.67it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "Extracting Dataset Version Zip to Banana-Ripeness-Classification-1 in folder:: 100%|██████████| 5639/5639 [00:00<00:00, 6735.45it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#Save the dataset name to the environment so we can use it in a system call later\n",
        "dataset_name = dataset.location.split(os.sep)[-1]\n",
        "os.environ[\"DATASET_NAME\"] = dataset_name"
      ],
      "metadata": {
        "id": "wLQbThFICpn4"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train On Custom Data 🎉\n",
        "Here, we use the DATASET_NAME environment variable to pass our dataset to the `--data` parameter.\n",
        "\n",
        "Note: we're training for 30 epochs here. We're also starting training from the pretrained weights. Larger datasets will likely benefit from longer training, frequently going as high as 1000-3000."
      ],
      "metadata": {
        "id": "-5z7Yv42FGrK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd ../yolov5\n",
        "!python classify/train.py --model yolov5s-cls.pt --data $DATASET_NAME --epochs 30 --img 128 --pretrained weights/yolov5s-cls.pt"
      ],
      "metadata": {
        "id": "MXWTTN2BEaqe",
        "outputId": "cdc90fce-4716-4550-becf-f6ef6f9a99de",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/yolov5\n",
            "2024-10-24 18:18:39.078290: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
            "2024-10-24 18:18:39.098773: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
            "2024-10-24 18:18:39.104998: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
            "\u001b[34m\u001b[1mclassify/train: \u001b[0mmodel=yolov5s-cls.pt, data=Banana-Ripeness-Classification-1, epochs=30, batch_size=64, imgsz=128, nosave=False, cache=None, device=, workers=8, project=runs/train-cls, name=exp, exist_ok=False, pretrained=weights/yolov5s-cls.pt, optimizer=Adam, lr0=0.001, decay=5e-05, label_smoothing=0.1, cutoff=None, dropout=None, verbose=False, seed=0, local_rank=-1\n",
            "\u001b[34m\u001b[1mgithub: \u001b[0mup to date with https://github.com/ultralytics/yolov5 ✅\n",
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n",
            "\n",
            "\u001b[34m\u001b[1mTensorBoard: \u001b[0mStart with 'tensorboard --logdir runs/train-cls', view at http://localhost:6006/\n",
            "\u001b[34m\u001b[1malbumentations: \u001b[0mRandomResizedCrop(p=1.0, size=(128, 128), scale=(0.08, 1.0), ratio=(0.75, 1.3333333333333333), interpolation=1), HorizontalFlip(p=0.5), ColorJitter(p=0.5, brightness=(0.6, 1.4), contrast=(0.6, 1.4), saturation=(0.6, 1.4), hue=(0.0, 0.0)), Normalize(p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, normalization='standard'), ToTensorV2(p=1.0, transpose_mask=False)\n",
            "Model summary: 149 layers, 4180166 parameters, 4180166 gradients, 10.5 GFLOPs\n",
            "\u001b[34m\u001b[1moptimizer:\u001b[0m Adam(lr=0.001) with parameter groups 32 weight(decay=0.0), 33 weight(decay=5e-05), 33 bias\n",
            "/content/yolov5/classify/train.py:201: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\n",
            "  scaler = amp.GradScaler(enabled=cuda)\n",
            "Image sizes 128 train, 128 test\n",
            "Using 8 dataloader workers\n",
            "Logging results to \u001b[1mruns/train-cls/exp\u001b[0m\n",
            "Starting yolov5s-cls.pt training on Banana-Ripeness-Classification-1 dataset with 6 classes for 30 epochs...\n",
            "\n",
            "     Epoch   GPU_mem  train_loss   test_loss    top1_acc    top5_acc\n",
            "  0% 0/62 [00:00<?, ?it/s]/content/yolov5/classify/train.py:222: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
            "  with amp.autocast(enabled=cuda):  # stability issues when enabled\n",
            "      1/30    0.531G        1.31                             testing:   0% 0/5 [00:00<?, ?it/s]/content/yolov5/classify/val.py:111: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
            "  with torch.cuda.amp.autocast(enabled=device.type != \"cpu\"):\n",
            "      1/30    0.531G        1.31        2.11       0.244       0.717: 100% 62/62 [00:04<00:00, 15.02it/s]\n",
            "      2/30    0.631G        1.08        1.12       0.598       0.996: 100% 62/62 [00:02<00:00, 23.08it/s]\n",
            "      3/30    0.631G        1.01        1.04       0.714           1: 100% 62/62 [00:02<00:00, 24.56it/s]\n",
            "      4/30    0.631G       0.942       0.868       0.779           1: 100% 62/62 [00:02<00:00, 25.07it/s]\n",
            "      5/30    0.631G       0.915       0.797       0.822           1: 100% 62/62 [00:02<00:00, 24.08it/s]\n",
            "      6/30    0.631G       0.857       0.747       0.854           1: 100% 62/62 [00:02<00:00, 21.40it/s]\n",
            "      7/30    0.631G       0.852       0.781       0.824           1: 100% 62/62 [00:02<00:00, 22.35it/s]\n",
            "      8/30    0.631G       0.827       0.795       0.836       0.998: 100% 62/62 [00:02<00:00, 24.24it/s]\n",
            "      9/30    0.631G       0.812       0.759       0.835           1: 100% 62/62 [00:02<00:00, 23.40it/s]\n",
            "     10/30    0.631G       0.776       0.716       0.865           1: 100% 62/62 [00:02<00:00, 23.94it/s]\n",
            "     11/30    0.631G       0.783       0.728       0.865           1: 100% 62/62 [00:02<00:00, 22.80it/s]\n",
            "     12/30    0.631G       0.774       0.653       0.906           1: 100% 62/62 [00:02<00:00, 21.32it/s]\n",
            "     13/30    0.631G       0.735       0.634        0.92           1: 100% 62/62 [00:02<00:00, 23.37it/s]\n",
            "     14/30    0.631G       0.735        0.64       0.916           1: 100% 62/62 [00:02<00:00, 23.80it/s]\n",
            "     15/30    0.631G       0.725       0.673       0.886           1: 100% 62/62 [00:02<00:00, 23.49it/s]\n",
            "     16/30    0.631G       0.704       0.648       0.879           1: 100% 62/62 [00:02<00:00, 24.17it/s]\n",
            "     17/30    0.631G       0.698       0.567       0.925           1: 100% 62/62 [00:02<00:00, 21.53it/s]\n",
            "     18/30    0.631G       0.664       0.558       0.938           1: 100% 62/62 [00:02<00:00, 21.99it/s]\n",
            "     19/30    0.631G       0.679       0.596        0.92           1: 100% 62/62 [00:02<00:00, 24.99it/s]\n",
            "     20/30    0.631G        0.65       0.587       0.927           1: 100% 62/62 [00:02<00:00, 24.38it/s]\n",
            "     21/30    0.631G       0.643       0.599       0.925           1: 100% 62/62 [00:02<00:00, 25.06it/s]\n",
            "     22/30    0.631G       0.636       0.547       0.941           1: 100% 62/62 [00:02<00:00, 23.08it/s]\n",
            "     23/30    0.631G       0.621       0.607       0.923           1: 100% 62/62 [00:02<00:00, 22.11it/s]\n",
            "     24/30    0.631G       0.636       0.567       0.931           1: 100% 62/62 [00:02<00:00, 22.56it/s]\n",
            "     25/30    0.631G       0.625        0.56       0.936           1: 100% 62/62 [00:02<00:00, 23.88it/s]\n",
            "     26/30    0.631G       0.608       0.544        0.94           1: 100% 62/62 [00:02<00:00, 25.20it/s]\n",
            "     27/30    0.631G         0.6       0.534       0.956           1: 100% 62/62 [00:02<00:00, 22.97it/s]\n",
            "     28/30    0.631G       0.585       0.526        0.95           1: 100% 62/62 [00:02<00:00, 21.88it/s]\n",
            "     29/30    0.631G       0.592       0.522        0.95           1: 100% 62/62 [00:02<00:00, 21.72it/s]\n",
            "     30/30    0.631G       0.585       0.524       0.954           1: 100% 62/62 [00:02<00:00, 23.64it/s]\n",
            "\n",
            "Training complete (0.024 hours)\n",
            "Results saved to \u001b[1mruns/train-cls/exp\u001b[0m\n",
            "Predict:         python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source im.jpg\n",
            "Validate:        python classify/val.py --weights runs/train-cls/exp/weights/best.pt --data /content/datasets/Banana-Ripeness-Classification-1\n",
            "Export:          python export.py --weights runs/train-cls/exp/weights/best.pt --include onnx\n",
            "PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'runs/train-cls/exp/weights/best.pt')\n",
            "Visualize:       https://netron.app\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Validate Your Custom Model\n",
        "\n",
        "Repeat step 2 from above to test and validate your custom model."
      ],
      "metadata": {
        "id": "HHUFGeLbGd98"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python classify/val.py --weights runs/train-cls/exp/weights/best.pt --data ../datasets/$DATASET_NAME"
      ],
      "metadata": {
        "id": "DIV7ydyKGZFL",
        "outputId": "a311138a-70f7-4b39-e81f-f14f4bad6a90",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[34m\u001b[1mclassify/val: \u001b[0mdata=../datasets/Banana-Ripeness-Classification-1, weights=['runs/train-cls/exp/weights/best.pt'], batch_size=128, imgsz=224, device=, workers=8, verbose=True, project=runs/val-cls, name=exp, exist_ok=False, half=False, dnn=False\n",
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n",
            "\n",
            "Fusing layers... \n",
            "Model summary: 117 layers, 4174374 parameters, 0 gradients, 10.4 GFLOPs\n",
            "testing:   0% 0/5 [00:00<?, ?it/s]/content/yolov5/classify/val.py:111: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
            "  with torch.cuda.amp.autocast(enabled=device.type != \"cpu\"):\n",
            "testing: 100% 5/5 [00:01<00:00,  2.70it/s]\n",
            "                   Class      Images    top1_acc    top5_acc\n",
            "                     all         562       0.929           1\n",
            "               freshripe         102        0.98           1\n",
            "             freshunripe          83           1           1\n",
            "                overripe         113       0.956           1\n",
            "                    ripe          52       0.731           1\n",
            "                  rotten         185       0.957           1\n",
            "                  unripe          27       0.593           1\n",
            "Speed: 0.1ms pre-process, 1.3ms inference, 0.2ms post-process per image at shape (1, 3, 224, 224)\n",
            "Results saved to \u001b[1mruns/val-cls/exp\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Infer With Your Custom Model"
      ],
      "metadata": {
        "id": "uH5tJNpEsi6g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Get the path of an image from the test or validation set\n",
        "if os.path.exists(os.path.join(dataset.location, \"test\")):\n",
        "  split_path = os.path.join(dataset.location, \"test\")\n",
        "else:\n",
        "  os.path.join(dataset.location, \"valid\")\n",
        "example_class = os.listdir(split_path)[0]\n",
        "example_image_name = os.listdir(os.path.join(split_path, example_class))[0]\n",
        "example_image_path = os.path.join(split_path, example_class, example_image_name)\n",
        "os.environ[\"TEST_IMAGE_PATH\"] = example_image_path\n",
        "\n",
        "print(f\"Inferring on an example of the class '{example_class}'\")\n",
        "\n",
        "# Infer\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source $TEST_IMAGE_PATH"
      ],
      "metadata": {
        "id": "81lK1hU_sk54",
        "outputId": "9868b2b9-cb1a-42df-fe0f-b6e4cc871a5a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Inferring on an example of the class 'unripe'\n",
            "\u001b[34m\u001b[1mclassify/predict: \u001b[0mweights=['runs/train-cls/exp/weights/best.pt'], source=/content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dcebc31-2653-11ec-874c-d8c4975e38aa-Copy_jpg.rf.7e4699ac9e8199f045f6f49829c42f94.jpg, data=data/coco128.yaml, imgsz=[224, 224], device=, view_img=False, save_txt=False, nosave=False, augment=False, visualize=False, update=False, project=runs/predict-cls, name=exp, exist_ok=False, half=False, dnn=False, vid_stride=1\n",
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n",
            "\n",
            "Fusing layers... \n",
            "Model summary: 117 layers, 4174374 parameters, 0 gradients, 10.4 GFLOPs\n",
            "image 1/1 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dcebc31-2653-11ec-874c-d8c4975e38aa-Copy_jpg.rf.7e4699ac9e8199f045f6f49829c42f94.jpg: 224x224 rotten 0.48, unripe 0.33, freshunripe 0.08, ripe 0.04, overripe 0.04, 3.4ms\n",
            "Speed: 0.3ms pre-process, 3.4ms inference, 30.7ms NMS per image at shape (1, 3, 224, 224)\n",
            "Results saved to \u001b[1mruns/predict-cls/exp3\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can see the inference results show ~3ms inference and the respective classes predicted probabilities."
      ],
      "metadata": {
        "id": "DdGuG-1kNjWT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## (OPTIONAL) Improve Our Model with Active Learning\n",
        "\n",
        "Now that we've trained our model once, we will want to continue to improve its performance. Improvement is largely dependent on improving our dataset.\n",
        "\n",
        "We can programmatically upload example failure images back to our custom dataset based on conditions (like seeing an underrpresented class or a low confidence score) using the same `pip` package."
      ],
      "metadata": {
        "id": "I38IM6NXKNN9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# # Upload example image\n",
        "# project.upload(image_path)\n"
      ],
      "metadata": {
        "id": "HycgSEnYKo0J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # Example upload code\n",
        "# min_conf = float(\"inf\")\n",
        "# for pred in results:\n",
        "#     if pred[\"score\"] < min_conf:\n",
        "#         min_conf = pred[\"score\"]\n",
        "# if min_conf < 0.4:\n",
        "#     project.upload(image_path)"
      ],
      "metadata": {
        "id": "VwXDoz_vLK3V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# (BONUS) YOLOv5 classify/predict.py Accepts Several Input Methods\n",
        "- Webcam: `python classify/predict.py --weights yolov5s-cls.pt --source 0`\n",
        "- Image `python classify/predict.py --weights yolov5s-cls.pt --source img.jpg`\n",
        "- Video: `python classify/predict.py --weights yolov5s-cls.pt --source vid.mp4`\n",
        "- Directory: `python classify/predict.py --weights yolov5s-cls.pt --source path/`\n",
        "- Glob: `python classify/predict.py --weights yolov5s-cls.pt --source 'path/*.jpg'`\n",
        "- YouTube: `python classify/predict.py --weights yolov5s-cls.pt --source 'https://youtu.be/Zgi9g1ksQHc'`\n",
        "- RTSP, RTMP, HTTP stream: `python classify/predict.py --weights yolov5s-cls.pt --source 'rtsp://example.com/media.mp4'`"
      ],
      "metadata": {
        "id": "aYlfaHDusN-j"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Directory Example"
      ],
      "metadata": {
        "id": "iKSP-SNTvcLJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Directory infer\n",
        "os.environ[\"TEST_CLASS_PATH\"] = test_class_path = os.path.join(*os.environ[\"TEST_IMAGE_PATH\"].split(os.sep)[:-1])\n",
        "print(f\"Infering on all images from the directory {os.environ['TEST_CLASS_PATH']}\")\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source /$TEST_CLASS_PATH/"
      ],
      "metadata": {
        "id": "lwSoHcHcvjeD",
        "outputId": "91ea2f7a-92b9-4e8d-a500-d212c76d82d5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Infering on all images from the directory content/datasets/Banana-Ripeness-Classification-1/test/unripe\n",
            "\u001b[34m\u001b[1mclassify/predict: \u001b[0mweights=['runs/train-cls/exp/weights/best.pt'], source=/content/datasets/Banana-Ripeness-Classification-1/test/unripe/, data=data/coco128.yaml, imgsz=[224, 224], device=, view_img=False, save_txt=False, nosave=False, augment=False, visualize=False, update=False, project=runs/predict-cls, name=exp, exist_ok=False, half=False, dnn=False, vid_stride=1\n",
            "YOLOv5 🚀 v7.0-378-g2f74455a Python-3.10.12 torch-2.5.0+cu121 CUDA:0 (NVIDIA A100-SXM4-40GB, 40514MiB)\n",
            "\n",
            "Fusing layers... \n",
            "Model summary: 117 layers, 4174374 parameters, 0 gradients, 10.4 GFLOPs\n",
            "image 1/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5202f0b6-2653-11ec-90a5-d8c4975e38aa-Copy_jpg.rf.b83969d5d68e6566a999d2001448efbe.jpg: 224x224 unripe 0.70, rotten 0.16, overripe 0.06, ripe 0.04, freshripe 0.03, 3.7ms\n",
            "image 2/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-53b71b36-2653-11ec-ac4b-d8c4975e38aa-Copy_jpg.rf.678ce536622a109f08931fda3e083df1.jpg: 224x224 unripe 0.72, rotten 0.18, overripe 0.04, ripe 0.03, freshripe 0.02, 5.7ms\n",
            "image 3/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-54487ec9-2653-11ec-b977-d8c4975e38aa-Copy_jpg.rf.67b4c28f725bd8ab3748124f393fc3a4.jpg: 224x224 unripe 0.68, rotten 0.18, overripe 0.06, ripe 0.04, freshripe 0.03, 4.7ms\n",
            "image 4/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-54b88418-2653-11ec-9666-d8c4975e38aa-Copy_jpg.rf.994892a288e6f9cc7177223dc59a4d60.jpg: 224x224 unripe 0.68, rotten 0.14, overripe 0.07, ripe 0.06, freshripe 0.03, 3.8ms\n",
            "image 5/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-54bae675-2653-11ec-a41d-d8c4975e38aa-Copy_jpg.rf.ee23804a7b2687e2fde689c2266c0c84.jpg: 224x224 unripe 0.75, rotten 0.12, overripe 0.05, ripe 0.05, freshripe 0.02, 3.7ms\n",
            "image 6/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5523c512-2653-11ec-a156-d8c4975e38aa-Copy_jpg.rf.a3c601a5df777fc73daf0a06e59ce2c9.jpg: 224x224 unripe 0.70, rotten 0.14, overripe 0.06, ripe 0.05, freshripe 0.03, 3.8ms\n",
            "image 7/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-55b52aee-2653-11ec-aa91-d8c4975e38aa-Copy_jpg.rf.de7a93a604a3f03598f15330b935999a.jpg: 224x224 unripe 0.65, rotten 0.15, overripe 0.08, ripe 0.06, freshripe 0.04, 3.7ms\n",
            "image 8/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-566f1326-2653-11ec-a976-d8c4975e38aa-Copy_jpg.rf.60c58821e931cae44f922a0665001685.jpg: 224x224 unripe 0.54, rotten 0.27, overripe 0.08, ripe 0.05, freshripe 0.03, 3.7ms\n",
            "image 9/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dc06f2e-2653-11ec-9e31-d8c4975e38aa-Copy_jpg.rf.d9626c3ef9c5fcf92042305014b37f92.jpg: 224x224 rotten 0.52, unripe 0.36, overripe 0.04, ripe 0.03, freshripe 0.03, 3.7ms\n",
            "image 10/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dcc59f1-2653-11ec-ac77-d8c4975e38aa-Copy_jpg.rf.f80272bad6c07d1ea217b67ab3716090.jpg: 224x224 unripe 0.43, rotten 0.40, overripe 0.06, ripe 0.05, freshunripe 0.04, 3.6ms\n",
            "image 11/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dcebc30-2653-11ec-908b-d8c4975e38aa-Copy_jpg.rf.321e01f4ef4e07b8e341d97d3b55ae89.jpg: 224x224 rotten 0.53, unripe 0.33, overripe 0.04, ripe 0.03, freshunripe 0.03, 3.6ms\n",
            "image 12/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5dcebc31-2653-11ec-874c-d8c4975e38aa-Copy_jpg.rf.7e4699ac9e8199f045f6f49829c42f94.jpg: 224x224 rotten 0.48, unripe 0.33, freshunripe 0.08, ripe 0.04, overripe 0.04, 3.7ms\n",
            "image 13/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5e032b73-2653-11ec-8c71-d8c4975e38aa-Copy_jpg.rf.313b94edc8a78238f3541264166d6086.jpg: 224x224 rotten 0.48, unripe 0.40, overripe 0.04, freshripe 0.03, ripe 0.03, 3.6ms\n",
            "image 14/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5e13da82-2653-11ec-93e8-d8c4975e38aa-Copy_jpg.rf.80c994a94ab4714e94da823ea4f0f207.jpg: 224x224 rotten 0.49, unripe 0.36, overripe 0.04, ripe 0.04, freshripe 0.04, 3.6ms\n",
            "image 15/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5e1fc54a-2653-11ec-a16c-d8c4975e38aa-Copy_jpg.rf.5b8202b68c3d8fee62f6371b031b4286.jpg: 224x224 rotten 0.50, unripe 0.34, ripe 0.06, overripe 0.04, freshripe 0.04, 3.7ms\n",
            "image 16/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5e817dbe-2653-11ec-8655-d8c4975e38aa-Copy_jpg.rf.dcdac1502165b8ad8d54336cc6af7e20.jpg: 224x224 unripe 0.43, rotten 0.41, ripe 0.05, freshunripe 0.04, freshripe 0.04, 3.6ms\n",
            "image 17/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5fadcdd9-2653-11ec-a703-d8c4975e38aa-Copy_jpg.rf.5dccbf617d6b522e341f07e676e3b3ac.jpg: 224x224 rotten 0.44, unripe 0.43, overripe 0.05, ripe 0.04, freshripe 0.03, 3.8ms\n",
            "image 18/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5fb02f80-2653-11ec-bcaa-d8c4975e38aa-Copy_jpg.rf.c1411afa8c7ca5cdae078f4729092082.jpg: 224x224 rotten 0.64, unripe 0.20, overripe 0.06, freshripe 0.05, ripe 0.04, 4.7ms\n",
            "image 19/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-5fb755f0-2653-11ec-b705-d8c4975e38aa-Copy_jpg.rf.9f0b62e72da0bb254a7c3db4756afb90.jpg: 224x224 unripe 0.48, rotten 0.35, overripe 0.06, ripe 0.05, freshripe 0.03, 3.8ms\n",
            "image 20/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-610e8c6e-2653-11ec-8930-d8c4975e38aa-Copy_jpg.rf.c354488439ff9d98d1624da8227db073.jpg: 224x224 rotten 0.45, unripe 0.42, overripe 0.05, freshripe 0.03, ripe 0.03, 3.8ms\n",
            "image 21/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-612b2644-2653-11ec-837f-d8c4975e38aa-Copy_jpg.rf.8f87af32f4b81399fdc3de36a5c0c0b1.jpg: 224x224 rotten 0.52, unripe 0.29, overripe 0.06, ripe 0.05, freshripe 0.05, 3.9ms\n",
            "image 22/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-6185b97e-2653-11ec-83e8-d8c4975e38aa-Copy_jpg.rf.e56b9060a63885cceaeb2123b1344c02.jpg: 224x224 unripe 0.61, rotten 0.27, overripe 0.04, freshunripe 0.04, ripe 0.02, 3.6ms\n",
            "image 23/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-61cf9b21-2653-11ec-a09d-d8c4975e38aa-Copy_jpg.rf.9e7a5273181b28eeda30f7d002e36f45.jpg: 224x224 unripe 0.51, rotten 0.37, overripe 0.04, freshripe 0.03, ripe 0.03, 3.6ms\n",
            "image 24/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-623879e6-2653-11ec-98ce-d8c4975e38aa_jpg.rf.19c13fde60aacdbc70cf655d4be5677d.jpg: 224x224 unripe 0.60, rotten 0.30, overripe 0.05, ripe 0.02, freshripe 0.02, 3.6ms\n",
            "image 25/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-7fecfc48-1d0a-11ec-b38a-d8c4975e38aa-1-_jpg.rf.ba3d265654f81f21a0f5ca611a709f34.jpg: 224x224 unripe 0.70, rotten 0.14, overripe 0.06, ripe 0.05, freshripe 0.03, 3.6ms\n",
            "image 26/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-86526c32-1d0a-11ec-9731-d8c4975e38aa_jpg.rf.49a628c187e9dd05b82373f9aba9ac2a.jpg: 224x224 rotten 0.70, unripe 0.17, overripe 0.05, ripe 0.04, freshripe 0.03, 3.6ms\n",
            "image 27/27 /content/datasets/Banana-Ripeness-Classification-1/test/unripe/musa-acuminata-unripe-87eec399-1d0a-11ec-8732-d8c4975e38aa-1-_jpg.rf.6c812af97e43756daed15e308b21d124.jpg: 224x224 unripe 0.65, rotten 0.25, overripe 0.04, ripe 0.02, freshripe 0.02, 3.8ms\n",
            "Speed: 0.3ms pre-process, 3.8ms inference, 1.2ms NMS per image at shape (1, 3, 224, 224)\n",
            "Results saved to \u001b[1mruns/predict-cls/exp4\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###YouTube Example"
      ],
      "metadata": {
        "id": "kCCao9t8se8i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#YouTube infer\n",
        "!python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source 'https://www.youtube.com/watch?v=7AlYA4ItA74'"
      ],
      "metadata": {
        "id": "heebjpJBsakV"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}